from collections.abc import Iterable, Iterator
from typing import Any, TypeVar, TypedDict

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, DatasetWriter, SkipExhausted, pause, \
    resume, skip
from zkl_aiutils_datasets.wrapping import EmptyDataset, EmptyDatasetIterator, EmptyDatasetWriter, wrap_dataset_iterator

AnySample = TypeVar('AnySample')


class ChainedIteratorState(TypedDict):
    children_iterator_state: Any
    samples_iterator_state: Any


class ChainedIterator(DatasetIterator[AnySample]):
    def __init__(self,
        children: Iterable[Iterable[AnySample]],
        state: ChainedIteratorState | None = None,
    ):
        if state is None:
            self._children_iterator_state_or_error = None, None
            self._children_iterator = wrap_dataset_iterator(children)
            self._samples_iterator = EmptyDatasetIterator()
        else:
            children_iterator_state = state['children_iterator_state']
            samples_iterator_state = state['samples_iterator_state']
            self._children_iterator_state_or_error = children_iterator_state, None
            self._children_iterator = resume(children, children_iterator_state)

            try:
                child = next(self._children_iterator)
            except StopIteration:
                child = EmptyDataset()
            self._samples_iterator = resume(child, samples_iterator_state)

    def __next__(self) -> AnySample:
        while True:
            try:
                return next(self._samples_iterator)
            except StopIteration:
                self._children_iterator_state_or_error = try_pause(self._children_iterator)

                try:
                    child = next(self._children_iterator)
                except StopIteration:
                    self._samples_iterator = EmptyDatasetIterator()
                    raise StopIteration
                self._samples_iterator = wrap_dataset_iterator(child)

    def __skip__(self, samples_n: int):
        if samples_n <= 0:
            return

        # skip samples from child
        try:
            skip(self._samples_iterator, samples_n)
            return
        except SkipExhausted as err:
            skipped_samples_n = err.skipped_samples_n

        # skip samples form children
        while skipped_samples_n < samples_n:
            self._children_iterator_state_or_error = try_pause(self._children_iterator)

            try:
                child = next(self._children_iterator)
            except StopIteration:
                self._samples_iterator = EmptyDatasetIterator()
                raise SkipExhausted(skipped_samples_n)

            try:
                # noinspection PyTypeChecker
                child_samples_n = len(child)
            except TypeError:
                child_samples_n = None

            if child_samples_n is not None:
                if child_samples_n == samples_n - skipped_samples_n:
                    self._children_iterator_state_or_error = try_pause(self._children_iterator)
                    self._samples_iterator = EmptyDatasetIterator()
                    return

                if child_samples_n < samples_n - skipped_samples_n:
                    skipped_samples_n += child_samples_n
                    continue

            self._samples_iterator = wrap_dataset_iterator(child)

            try:
                skip(self._samples_iterator, samples_n - skipped_samples_n)
                return
            except SkipExhausted as err:
                skipped_samples_n += err.skipped_samples_n

    def __pause__(self) -> ChainedIteratorState | None:
        children_iterator_state, children_iterator_error = self._children_iterator_state_or_error
        if children_iterator_error is not None:
            raise children_iterator_error
        samples_iterator_state = pause(self._samples_iterator)
        return ChainedIteratorState(
            children_iterator_state=children_iterator_state,
            samples_iterator_state=samples_iterator_state)


class ChainedDatasets(Dataset[AnySample]):
    """ A dataset that produces samples from a (potentially large) sequence of sub-datasets in order. """

    def __init__(self, children: Iterable[Iterable[AnySample]]):
        self._children = children

    @property
    def children(self) -> Iterable[Iterable[AnySample]]:
        return self._children

    def __iter__(self) -> ChainedIterator[AnySample]:
        return ChainedIterator(self._children)

    def __resume__(self, state: ChainedIteratorState | None = None, /) -> ChainedIterator[AnySample]:
        return ChainedIterator(self._children, state)


class ChainedDatasetsWriter(DatasetWriter[AnySample]):
    def __init__(self, children: Iterable[DatasetWriter[AnySample]] | Iterator[DatasetWriter[AnySample]]):
        self._children_iterator = iter(children)
        self._child = EmptyDatasetWriter()

    def write(self, sample: AnySample):
        while True:
            try:
                self._child.write(sample)
                break
            except StopIteration:
                self._child.close()
                self._child = next(self._children_iterator)

    def flush(self):
        self._child.flush()

    def close(self):
        self._child.close()


# utils

def try_pause(iterator: Iterator[Any]):
    try:
        return pause(iterator), None
    except TypeError as error:
        return None, error
